import torch
import torch.nn.functional as F
from train import LanguageModel, tokenizer, embedding, device
from sys3_位置编码 import PositionEncoding

# 加载模型
model = LanguageModel().to(device)
model.load_state_dict(torch.load("weights/best_model.pth", map_location=device))
model.eval()
print("✅ 模型加载完成")

# ========== 单轮问答 ==========
def ask(question, context=None, max_len=100, temperature=1.0, top_k=50):
    """
    question: 用户输入的问题
    context: 选填，多轮上下文，可以是字符串或列表
    temperature: 控制生成文本的随机性，值越低越确定，越高越随机
    top_k: Top-K 采样参数，只从概率最高的 K 个词中采样
    """
    # 构建上下文文本，格式与训练时保持一致
    context_parts = []
    
    # 添加系统提示，与训练数据格式保持一致
    system_prompt = "你是一位精通理情行为疗法（Rational Emotive Behavior Therapy，简称REBT）的心理咨询师，能够合理地采用理情行为疗法给来访者提供专业地指导和支持，缓解来访者的焦虑、抑郁等负面情绪，并简要描述了各个阶段的重点。\n（1）**检查非理性信念和自我挫败式思维**：理情行为疗法把认知干预视为治疗的“生命”，因此，几乎从治疗一开始，在问题探索阶段，咨询师就以积极的、说服教导的方式，引导来访者去觉察自己在感到焦虑、抑郁或愤怒前对自己“说”了些什么。\n（2）**与非理性信念辩论**：咨询师运用多种技术（主要是认知技术）帮助来访者向非理性信念和思维质疑发难，证明它们的不现实、不合理之处，激活其内在意愿，放弃非理性信念，建立合理信念，找出理性的情感和事件的适宜的、理性的反应，找出理性的信念和实事求是的、指向问题解决的思维陈述，以此来替代非理性信念和自我挫败式思维。为了巩固理性信念，咨询师要向来访者反复教导，证明为什么理性信念和思维比非理性信念和思维更符合客观现实的态度，科学合理的思维方式内化成个人的生活态度，并在以后的生活中坚持不懈地按理情行为疗法的教导来解决新的问题。"
    context_parts.append(system_prompt)
    
    if context:
        if isinstance(context, list):
            for i, c in enumerate(context):
                prefix = "用户：" if i % 2 == 0 else "心理师："
                context_parts.append(prefix + c)
        else:
            context_parts.append(context)
    
    # 添加当前问题，格式与训练时保持一致
    context_parts.append("用户：" + question)
    # 添加心理师的标识，模型将从这里开始生成回复
    context_parts.append("心理师：")
    
    src_text = "\n".join(context_parts)
    # print(f"DEBUG: 输入文本为: {repr(src_text)}")  # 调试信息
    
    src_ids, src_mask = tokenizer([src_text])
    src_ids, src_mask = src_ids.to(device), src_mask.to(device)

    # 起始 token
    start_token = torch.tensor([[embedding.tokenizer.cls_token_id or 101]], device=device)
    tgt = start_token
    tgt_mask = torch.zeros((1, 1), dtype=torch.bool, device=device)

    for _ in range(max_len):
        with torch.no_grad():
            out = model(src_ids, tgt, src_mask, tgt_mask)
            next_token_logits = out[:, -1, :] / temperature  # 应用温度系数
            
            # Top-K 采样
            if top_k > 0:
                top_k = min(top_k, next_token_logits.size(-1))  # 确保不超过词汇表大小
                # 获取top-k最高的logits索引
                top_k_logits, top_k_indices = torch.topk(next_token_logits, top_k, dim=-1)
                # 将非top-k的logits设为负无穷大
                next_token_logits_filtered = torch.full_like(next_token_logits, float('-inf'))
                next_token_logits_filtered.scatter_(-1, top_k_indices, top_k_logits)
            
            # 计算概率分布
            probs = F.softmax(next_token_logits, dim=-1)
            
            # 从概率分布中采样
            next_token = torch.multinomial(probs, 1)

        tgt = torch.cat([tgt, next_token], dim=1)
        tgt_mask = torch.zeros((1, tgt.size(1)), dtype=torch.bool, device=device)

        if next_token.item() == embedding.tokenizer.sep_token_id:
            break

    reply = embedding.tokenizer.decode(tgt[0].tolist(), skip_special_tokens=True)
    # 处理多余空格问题
    reply = " ".join(reply.split())  # 使用空格连接而不是完全删除
    return reply


# ========== 运行示例 ==========
if __name__ == "__main__":
    print("🧠 心理咨询问答系统已启动\n")
    context = []  # 用于存储对话历史

    while True:
        question = input("👤 用户：").strip()
        if question in ["退出", "exit", "quit"]:
            print("👋 再见！")
            break

        # 使用更合适的参数提高回答相关性
        reply = ask(question, context, temperature=0.8, top_k=0)
        print("🩵 心理师：", reply)

        # 更新上下文，确保正确维护对话历史
        context.extend([question, reply])
        # 限制上下文长度，避免过长
        if len(context) > 10:  # 保持最近5轮对话
            context = context[-10:]